from langchain_core.example_selectors.semantic_similarity import (
    SemanticSimilarityExampleSelector,
)

from langchain_chroma import Chroma
from open_ai import get_open_ai_embeddings

from langchain_core.prompts.few_shot import FewShotPromptTemplate
from langchain_core.prompts.prompt import PromptTemplate

examples = [
    {
        "question": "谁的寿命更长，张三还是李四？",
        "answer": """
            这里需要跟进问题吗：是的。
            跟进：张三去世时多大？
            中间答案：张三去世时74岁。
            跟进：李四去世时多大？
            中间答案：李四去世时41岁。
            所以最终答案是：张三。
            """,
    },
    {
        "question": "craigslist的创始人是什么时候出生的？",
        "answer": """
            这里需要跟进问题吗：是的。
            跟进：craigslist的创始人是谁？
            中间答案：craigslist由Crag Newmark创立。
            跟进：Crag Newmark是什么时候出生的？
            中间答案：Crag Newmark于1952年12月6日出生。
            所以最终答案是：1952年12月6日
            """,
    },
    {
        "question": "《大白鲨》和《皇家赌场》的导演都来自同一个国家吗？",
        "answer": """
            这里需要跟进问题吗：是的。
            跟进：《大白鲨》的导演是谁？
            中间答案：《大白鲨》的导演是Steven Spielberg。
            跟进：Steven Spielberg来自哪里？
            中间答案：美国。
            跟进：《皇家赌场》的导演是谁？
            中间答案：《皇家赌场》的导演是Martin Campbell。
            跟进：Martin Campbell来自哪里？
            中间答案：新西兰。
            所以最终答案是：不是
            """,
    },
]

# 使用语义相似性示例选择器
example_selector = SemanticSimilarityExampleSelector.from_examples(
    # 这是可供选择的示例列表
    examples,
    # 这是用于生成嵌入的嵌入类，该嵌入用于衡量语义相似性 from langchain_openai import OpenAIEmbeddings
    get_open_ai_embeddings(),
    # 这是用于存储嵌入和执行相似性搜索的VectorStore类
    Chroma,
    # 这是要生成的示例数
    k=1,
)

question = "《大白鲨》导演是哪个国家的？"
selected_examples = example_selector.select_examples({"question": question})
print(f"最相似的示例：{question}")
for example in selected_examples:
    print("\n")
    for k, v in example.items():
        print(f"{k}:{v}")


example_prompt = PromptTemplate(
    input_variables=["question", "answer"], template="问题：{question}\\n{answer}"
)
prompt = FewShotPromptTemplate(
    example_selector=example_selector,
    example_prompt=example_prompt,
    suffix="问题：{input}",
    input_variables=["input"],
)
print(prompt.format(input="《大白鲨》导演是哪个国家的？"))
